{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 54,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "已加载标签到ID的映射字典：\n",
      "{'2': 0, '1': 1, '4': 2, '5': 3, '3': 4}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "# 从JSON文件中加载标签到ID的映射字典\n",
    "with open('label2id.json', 'r', encoding='utf-8') as f:\n",
    "    label2id_dict = json.load(f)\n",
    "\n",
    "print(\"已加载标签到ID的映射字典：\")\n",
    "print(label2id_dict)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-20T07:51:58.532501700Z",
     "start_time": "2024-06-20T07:51:58.516780600Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "class TempDataset(Dataset):\n",
    "    def __init__(self, texts, tokenizer, label2id, max_length):\n",
    "        self.texts = texts\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "        self.label2id = label2id\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.texts)\n",
    "    \n",
    "    def find_split_point(self, text, start, max_length):\n",
    "        split_point = min(start + max_length, len(text) - 1)\n",
    "        while split_point > start:\n",
    "            if text[split_point] in \"。！？：；\":\n",
    "                if split_point - start + 1 <= max_length:\n",
    "                    return split_point + 1\n",
    "            split_point -= 1\n",
    "        return start + max_length\n",
    "    \n",
    "    def segment_and_tokenize(self, text):\n",
    "        input_ids, attention_masks = [], []\n",
    "        i = 0\n",
    "        while i < len(text):\n",
    "            if len(text[i:]) > self.max_length - 2:\n",
    "                split_point = self.find_split_point(text, i, self.max_length - 2)\n",
    "            else:\n",
    "                split_point = len(text)\n",
    "            segment_text = text[i:split_point]\n",
    "            encoded_dict = self.tokenizer.encode_plus(\n",
    "                segment_text,\n",
    "                is_split_into_words=True,\n",
    "                max_length=self.max_length,\n",
    "                padding='max_length',\n",
    "                truncation=True,\n",
    "                return_attention_mask=True\n",
    "            )\n",
    "            input_ids.append(encoded_dict['input_ids'])\n",
    "            attention_masks.append(encoded_dict['attention_mask'])\n",
    "            i = split_point\n",
    "    \n",
    "        input_ids = torch.tensor(input_ids, dtype=torch.long)\n",
    "        attention_masks = torch.tensor(attention_masks, dtype=torch.long)\n",
    "\n",
    "        return {\n",
    "            'input_ids': input_ids,\n",
    "            'attention_mask': attention_masks\n",
    "        }\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        text = self.texts[idx]\n",
    "        return self.segment_and_tokenize(text)\n",
    "    \n",
    "    def collate_fn(self, batch):\n",
    "        input_ids = torch.cat([item['input_ids'] for item in batch], dim=0)\n",
    "        attention_mask = torch.cat([item['attention_mask'] for item in batch], dim=0)\n",
    "\n",
    "        return {\n",
    "            'input_ids': input_ids,\n",
    "            'attention_mask': attention_mask\n",
    "        }\n",
    "    "
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-20T07:52:00.682869900Z",
     "start_time": "2024-06-20T07:52:00.661402400Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['成分', '味道', '礼品', '快递小哥']\n",
      "[(0, 1), (40, 41), (63, 64), (72, 75)]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from keywordpredic import kwinference\n",
    "from transformers import BertTokenizer, ErnieForSequenceClassification\n",
    "\n",
    "predictor = kwinference(needEmbed=False)\n",
    "texts = '成分很安全的一款面膜，对香精敏感的美眉慎入，作为日常补水款足够了。滋润，不油腻，味道也喜欢。一直在用，感觉还是不错，这次还有小礼品，谢谢啦！还有快递小哥，今天下了雪如期送达，感谢！不敢相信这么便宜的面膜一点不比八块十块的差。面膜纸丝薄服帖我很喜欢，补水用也不会心疼。回回购'\n",
    "attr,view,positions= predictor(texts)\n",
    "print(attr)\n",
    "print(positions)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-20T07:52:04.220555600Z",
     "start_time": "2024-06-20T07:52:01.766898300Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['成分很安全的一款面膜，对香精敏感的美眉慎入，作为日常补水款足够了。', '滋润，不油腻，味道也喜欢。', '一直在用，感觉还是不错，这次还有小礼品，谢谢啦！', '还有快递小哥，今天下了雪如期送达，感谢！']\n"
     ]
    }
   ],
   "source": [
    "def find_sentence_boundaries(text, start, end):\n",
    "    # 向前寻找句子的起始边界\n",
    "    while start > 0 and text[start - 1] not in \"。！？：；\":\n",
    "        start -= 1\n",
    "\n",
    "    # 向后寻找句子的结束边界\n",
    "    while end < len(text) - 1 and text[end + 1] not in \"。！？：；\":\n",
    "        end += 1\n",
    "\n",
    "    # 包含最后的标点符号\n",
    "    if end < len(text) - 1:\n",
    "        end += 1\n",
    "\n",
    "    return start, end\n",
    "\n",
    "def split_texts_by_positions(text, positions):\n",
    "    segments = []\n",
    "    for start, end in positions:\n",
    "        # 找到包含位置的完整句子\n",
    "        sentence_start, sentence_end = find_sentence_boundaries(text, start, end)\n",
    "        segments.append(text[sentence_start:sentence_end + 1])\n",
    "    \n",
    "    return segments\n",
    "\n",
    "# 分割文本\n",
    "segments = split_texts_by_positions(texts, positions)\n",
    "print(segments)\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-20T07:52:06.203278600Z",
     "start_time": "2024-06-20T07:52:06.174070400Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of ErnieForSequenceClassification were not initialized from the model checkpoint at ../nlp/ernie-3.0-mini-zh and are newly initialized: ['classifier.bias', 'classifier.weight', 'ernie.pooler.dense.bias', 'ernie.pooler.dense.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": "{'input_ids': tensor([[    1,    33,    59,   321,   204,    62,     5,     7,   699,    76,\n           1154,     4,    51,   673,   326,  1443,   345,     5,   188,  2116,\n           2143,   109,     4,    25,    13,   139,   223,   807,   101,   699,\n            581,   824,    15, 12043,     2,     0,     0,     0,     0,     0,\n              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n              0,     0,     0,     0,     0,     0,     0,     0]]),\n 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n          0, 0, 0, 0, 0, 0, 0, 0]])}"
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 加载预训练的ERNIE模型和tokenizer\n",
    "model_name = \"../nlp/ernie-3.0-mini-zh\" \n",
    "tokenizer = BertTokenizer.from_pretrained(model_name)\n",
    "model = ErnieForSequenceClassification.from_pretrained(model_name,num_labels=len(label2id_dict))\n",
    "model.load_state_dict(torch.load('model.pth',map_location=torch.device('cpu')))\n",
    "# 创建临时数据集和 DataLoader\n",
    "temp_dataset = TempDataset(segments,tokenizer, label2id_dict, max_length=128)\n",
    "temp_loader = DataLoader(temp_dataset, batch_size=2, collate_fn=temp_dataset.collate_fn)\n",
    "temp_dataset[0]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-20T07:52:07.593701400Z",
     "start_time": "2024-06-20T07:52:07.337402200Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "成分很安全的一款面膜，对香精敏感的美眉慎入，作为日常补水款足够了。->成分 :3\n",
      "滋润，不油腻，味道也喜欢。->味道 :3\n",
      "一直在用，感觉还是不错，这次还有小礼品，谢谢啦！->礼品 :3\n",
      "还有快递小哥，今天下了雪如期送达，感谢！->快递小哥 :3\n"
     ]
    }
   ],
   "source": [
    "predictions = []\n",
    "with torch.no_grad():\n",
    "    for batch in temp_loader:\n",
    "        input_ids = batch['input_ids']\n",
    "        attention_mask = batch['attention_mask']\n",
    "        outputs = model(input_ids, attention_mask=attention_mask)\n",
    "        logits = outputs.logits\n",
    "        preds = torch.argmax(logits, dim=-1)\n",
    "        predictions.extend(preds.cpu().numpy())\n",
    "\n",
    "id2label = {v: k for k, v in label2id_dict.items()}\n",
    "predicted_labels = [id2label[pred] for pred in predictions]\n",
    "for i in range(len(segments)):\n",
    "    print(f'{segments[i]}->{attr[i]} :{predicted_labels[i]}')"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-20T08:04:53.615580400Z",
     "start_time": "2024-06-20T08:04:53.380457400Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "huggingface",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
